from time import time
import matplotlib
import matplotlib.pyplot as plt
# matplotlib.use('Qt5Agg')
import pandas as pd
import numpy as np
import math
import cv2

from myUtils.curvefitting import CurveFit
from myUtils.kalmantrack import KalmanTrack
from myUtils.resultingdata import ResultingData
from myUtils.utils import xywh_to_xyxy, cal_iou


class QueueInformationDetection:
    IOU_Threshold = 0.3

    def __init__(self):
        self.resulting_data = ResultingData()
        # 队伍有多少人
        self.people_in_line = 0
        # 队伍的速度
        self.line_speed_find_target_state = True
        self.tracker_plan = False
        self.left_person_n1 = 0
        self.right_person_n2 = 0
        self.left_person_location = []
        self.right_person_location = []
        self.line_speed = 0
        self.target_person = []
        self.t0 = 0

        self.pre_target = []
        self.stay_count = 0

    def get_queue_info(self, image, number_of_person):
        self.queue_image = image
        self.number_of_persons = int(number_of_person)
        self.persons_position_list = []
        self.center_positions = []
        self.error_points = []
        self.max_iou_matched = False

    def get_persons_position_list(self, position_bbox):
        x0 = int(position_bbox[0])
        y0 = int(position_bbox[1])
        x1 = int(position_bbox[2])
        y1 = int(position_bbox[3])
        one_position = [x0, y0, x1, y1]
        self.persons_position_list.append(one_position)

    def set_persons_position_list(self, position_list):
        self.persons_position_list = position_list

    def get_persons_number(self):
        # test，加入异常值测试用------------------------------------------
        # self.persons_position_list.append([4, 4, 8, 8])
        # self.persons_position_list.append([80, 4, 64, 8])
        # -------------------------------------------------------------

        # 1. 计算所有矩形框上边界的坐标
        for position in self.persons_position_list:
            center_x = int((position[2] - position[0]) / 2) + position[0]
            center_y = int((position[3] - position[1]) / 2) + position[1]
            self.center_positions.append([center_x, center_y])

        # 2. 可视化所有点
        # test，在原图中可视化------------------------------------------
        # circle_img = self.queue_image.copy()
        # for cp in self.center_positions:
        #     cp = tuple(cp)
        #     cv2.circle(circle_img, cp, 5, (0, 255, 255), thickness=10)
        #
        # cv2.imshow('circle_img', circle_img)
        # cv2.waitKey()
        # cv2.imwrite(r'E:\2022研电赛\code\yolov5-master\line_detector\scripts\images\a.jpg', circle_img)
        # ------------------------------------------------------------
        center_array_row = np.array(self.center_positions)
        # test, 所有值可视化----------------------------------------------
        # fig = plt.figure()
        # ax1 = fig.add_subplot(111)
        # ax1.scatter(center_array_row[:, 0], center_array_row[:, 1], color='blue')
        # fig.show()
        # ------------------------------------------------------------

        if len(self.persons_position_list) <= 4:
            self.people_in_line = len(self.persons_position_list)
            self.resulting_data.get_people_in_line(self.people_in_line)
            # print('people: ' + str(self.people_in_line))
            return self.people_in_line
        # 3. 数据拟合 - 最小二乘拟合
        # 3.1. 剔除图像中的坏值
        # 按x方向从大到小排序
        self.center_positions_ordered = sorted(self.center_positions, key=(lambda x: x[0]), reverse=False)
        # 前后两个中心点的距离
        dis_center = []
        for i in range(0, len(self.center_positions_ordered) - 1):
            dis_center.append(math.sqrt((self.center_positions_ordered[i + 1][0] - self.center_positions_ordered[i][0]) *
                                        (self.center_positions_ordered[i + 1][0] - self.center_positions_ordered[i][0]) +
                                        (self.center_positions_ordered[i + 1][1] - self.center_positions_ordered[i][1]) *
                                        (self.center_positions_ordered[i + 1][1] - self.center_positions_ordered[i][1])))

        # 用箱型图检测是否存在异常值 - 计算出异常值
        fig2 = plt.figure(figsize=(10, 8))
        ax_tmp = fig2.add_subplot(1, 2, 1)
        bp = ax_tmp.boxplot(dis_center)

        lower_whisker = [item.get_ydata()[1] for item in bp['whiskers']][0]  # 上边缘
        upper_whisker = [item.get_ydata()[1] for item in bp['whiskers']][1]  # 下边缘
        # print("上边缘：", upper_whisker)
        # print("下边缘：", lower_whisker)
        # print("非异常范围：", [lower_whisker, upper_whisker])

        dis_center = np.array(dis_center)
        outliers = dis_center[(dis_center < lower_whisker) | (dis_center > upper_whisker)]
        # print("异常值（离散值）：", outliers)  #  [487.54 488.89 476.38] ndarray

        # 去除异常值
        clear_center_position = self.center_positions_ordered.copy()  # 去掉异常值的数组
        self.drop_points = []  # 去掉点的坐标

        for outlier in outliers:
            outlier_index = np.argwhere(dis_center == outlier)  # 异常值坐标
            # print(outlier_index[0][0])  # 如果坐标是i和i + 1都有可能有问题
            idx1 = outlier_index[0][0]
            idx2 = idx1 + 1

            if clear_center_position[idx1] != 0:
                self.drop_points.append(self.center_positions_ordered[idx1])
                clear_center_position[idx1] = 0

            if clear_center_position[idx2] != 0:
                self.drop_points.append(self.center_positions_ordered[idx2])
                clear_center_position[idx2] = 0

        j = 0
        for i in range(len(clear_center_position)):
            if clear_center_position[j] == 0:
                clear_center_position.pop(j)
            else:
                j += 1

        # 去除异常值后的数组，如果小于等于4也是无法进行拟合
        if len(clear_center_position) <= 4:
            self.people_in_line = len(clear_center_position)
            self.resulting_data.get_people_in_line(self.people_in_line)
            # print('people: ' + str(self.people_in_line))
            return self.people_in_line

        # 3.2. 曲线进行拟合（最小二乘）
        center_positions_array = np.array(clear_center_position)
        center_positions_array_x = center_positions_array[:, 0]
        center_positions_array_y = center_positions_array[:, 1]

        # 尝试1、2、3次项拟合，取残差最小
        curve_fitting = CurveFit(center_positions_array)
        final_fit_curve = center_positions_array
        residual = float('inf')
        for i in range(1, 4):
            # 曲线拟合
            curve = curve_fitting.fit_process(i)
            final_fit_curve = curve

            # 计算残差
            res = curve_fitting.calculate_residual()

            if res < residual:
                residual = res
                final_fit_curve = curve


        self.right_points, self.error_points = curve_fitting.outliers_re_judgment(self.drop_points)
        self.people_in_line = len(clear_center_position) + len(self.right_points)
        self.resulting_data.get_people_in_line(self.people_in_line)
        # print('people: ' + str(self.people_in_line))

        # test, 绘图查看----------------------------------------------
        # fig = plt.figure()
        # ax1 = fig.add_subplot(111)
        # # 图中所有人的散点图
        # ax1.scatter(center_array_row[:, 0], center_array_row[:, 1], color='blue', label='all points')
        # # 初步判断在队伍中的人的散点图
        # ax1.scatter(center_positions_array_x, center_positions_array_y, color='red', label='fit points')
        # # 加入误判的点
        # right_points_array = np.array(self.right_points)
        # ax1.scatter(right_points_array[:, 0], right_points_array[:, 1], color='purple', label='right points')
        # # 在队伍中人的拟合曲线
        # ax1.plot(center_positions_array_x, final_fit_curve, linewidth=3, color='g', label='fit curve')
        #
        # all_fit_curve = curve_fitting.complete_curve(np.array(self.center_positions_ordered))
        # # 加上误判后的人的拟合曲线
        # ax1.plot(np.array(self.center_positions_ordered)[:, 0], all_fit_curve, color='yellow', label='final curve')
        # ax1.legend(loc='best')
        # fig.show()
        # ------------------------------------------------------------

        return self.people_in_line

    def max_match(self, now_person):
        max_iou = self.IOU_Threshold
        matched_box = []
        for j, data_ in enumerate(self.person_data_sorted['person_position']):
            data_ = list(data_)
            iou = cal_iou(now_person, data_)
            if iou > max_iou:
                matched_box = data_
                max_iou = iou
                self.max_iou_matched = True
        return matched_box

    def get_speed_of_the_queue(self, plot_mark, fps):
        # 存在两个状态，一个是在跟踪状态/寻找目标状态

        # 1. 选定队伍中间人的位置
        self.person_data = pd.DataFrame({'index': [i for i in range(0, len(self.persons_position_list))],
                                         'person_position': self.persons_position_list,
                                         'person_center': self.center_positions})
        # 去除异常点
        for ep in self.error_points:
            self.person_data.drop(self.person_data[self.person_data['person_center'].map(tuple) == tuple(ep)].index,
                                  inplace=True)
        # 对在队伍中的人进行排序
        self.person_data_sorted = self.person_data.sort_values(by='person_center', ascending=True, ignore_index=True)

        if self.line_speed_find_target_state:
            self.kalman_tracker = KalmanTrack()
            # 取队伍中间的人为目标任务
            self.target_person = self.person_data_sorted.loc[int(len(self.person_data_sorted['index']) / 2)]
            self.target_person = self.target_person['person_position']

            # test, 目标人物原图可视化-------------------------------------------
            # target_img = self.queue_image.copy()
            # cv2.rectangle(target_img, (target_person[0], target_person[1]),
            #               (target_person[2], target_person[3]), (0, 255, 255), thickness=5)
            # cv2.imshow('target_img', target_img)
            # cv2.waitKey()
            # ------------------------------------------------------------
            # 在原图中进行标记
            if plot_mark:
                cv2.rectangle(self.queue_image, (self.target_person[0], self.target_person[1]),
                              (self.target_person[2], self.target_person[3]), (0, 255, 255), thickness=2)

            if int(len(self.person_data_sorted['index']) / 2) + 1 > len(self.person_data_sorted['index']) - 1 or \
                int(len(self.person_data_sorted['index']) / 2) - 1 < 0:
                # 方案1：计算当前人前面和后面各有几个人n1 和 n2,记录当前的时间戳t0
                self.tracker_plan = True
                self.left_person_n1 = int(len(self.person_data_sorted['index']) / 2)
                self.right_person_n2 = len(self.person_data_sorted['index']) - self.left_person_n1 - 1
            else:
                self.tracker_plan = False
                self.left_person_location = self.person_data_sorted.loc[
                    int(len(self.person_data_sorted['index']) / 2) - 1]
                self.left_person_location = self.left_person_location['person_position']

                self.right_person_location = self.person_data_sorted.loc[
                    int(len(self.person_data_sorted['index']) / 2) + 1]
                self.right_person_location = self.right_person_location['person_position']

            self.t0 = time()  # 返回当前时间的时间戳（1970纪元后经过的浮点秒数）。

            # 卡尔曼滤波器放入初值
            self.kalman_tracker.get_initial_state(self.target_person)
            self.line_speed_find_target_state = False

        else:
            if len(self.target_person) == 0:
                return 0

            # 对目标进行跟踪
            self.target_person = self.max_match(self.target_person)

            # test----------------------------------
            # self.target_person = self.persons_position_list[0]
            # --------------------------------------
            self.target_person = self.kalman_tracker.correct_and_predict(self.target_person, self.max_iou_matched)

            # 转换成list
            self.target_person = [self.target_person[0][0], self.target_person[1][0],
                                  self.target_person[2][0], self.target_person[3][0]]

            if self.pre_target == self.target_person:
                self.stay_count += 1

            if self.stay_count >= 5 * fps:
                self.line_speed_find_target_state = True
                self.stay_count = 0

            self.pre_target = self.target_person

            if plot_mark:
                self.target_person = xywh_to_xyxy(self.target_person)
                cv2.rectangle(self.queue_image, (self.target_person[0], self.target_person[1]),
                              (self.target_person[2], self.target_person[3]), (0, 255, 255), thickness=2)

            x_limit = self.queue_image.shape[1]
            y_limit = self.queue_image.shape[0]



            # 对于跟踪失败跟踪框乱飞，那么需要重新寻找目标人物
            if self.target_person[1] <= 0 and self.target_person[2] < x_limit and self.target_person[0] > 0 or \
                    self.target_person[3] >= y_limit and self.target_person[2] < x_limit and self.target_person[0] > 0:
                self.line_speed_find_target_state = True
                return self.line_speed

            if self.tracker_plan:
                if self.target_person[2] >= x_limit or self.target_person[0] <= 0:
                    t1 = time()
                    if self.target_person[2] == x_limit:
                        self.line_speed = (t1 - self.t0) / (self.right_person_n2 + 1)
                    else:
                        self.line_speed = (t1 - self.t0) / (self.left_person_n1 + 1)
                    if self.line_speed > 1:
                        self.line_speed_find_target_state = True
                        if fps:
                            self.line_speed = self.line_speed * 25 / fps
                        else:
                            self.line_speed = 0
                        self.resulting_data.get_line_speed(self.line_speed)
                # print('speed: ' + str(self.line_speed))
            else:
                # print(self.left_person_location)
                # print(self.right_person_location)
                if self.target_person[2] >= self.right_person_location[2] or \
                        self.target_person[2] <= self.left_person_location[2]:
                    t1 = time()
                    self.line_speed = (t1 - self.t0)

                    if self.line_speed > 1:
                        self.line_speed_find_target_state = True
                        if fps:
                            self.line_speed = self.line_speed * 25 / fps
                        else:
                            self.line_speed = 0
                        self.resulting_data.get_line_speed(self.line_speed)
            return self.line_speed

    def return_result(self):
        return self.resulting_data


if __name__ == '__main__':
    img = cv2.imread(r'../data/images/102865217.jpg')
    q = QueueInformationDetection(img, 33)
    position_list_test = [[1311, 463, 1377, 683], [1293, 467, 1345, 685], [209, 486, 263, 681], [775, 457, 833, 705], [249, 485, 299, 688], [1481, 468, 1535, 674], [1569, 474, 1600, 665], [1499, 471, 1560, 674], [1068, 461, 1137, 696], [106, 487, 158, 683], [7, 491, 52, 682], [1399, 470, 1452, 683], [279, 480, 322, 685], [1437, 464, 1493, 679], [477, 478, 530, 689], [665, 464, 723, 695], [160, 480, 225, 688], [367, 478, 425, 691], [1265, 474, 1329, 686], [307, 468, 370, 687], [1337, 469, 1400, 684], [1137, 455, 1200, 700], [872, 450, 940, 705], [61, 480, 121, 683], [935, 450, 1029, 708], [563, 464, 635, 699], [511, 470, 583, 694], [414, 479, 487, 691], [719, 450, 787, 705], [615, 456, 677, 702], [1199, 465, 1271, 697], [814, 461, 885, 707], [1023, 457, 1088, 715]]
    q.set_persons_position_list(position_list_test)
    q.get_persons_number()
    q.get_speed_of_the_queue()

